// import { ImageAsset, Texture2D } from "cc";
import { TensorShape } from "../../Core/Tensor";
import { CompressionSpec, SensorCompressionType } from "./CompressionSpec";
import { BuiltInSensorType, IBuiltInSensor } from "./IBuiltInSensor";
import { ISensor, SensorExtensions } from "./ISensor";
import { ObservationSpec } from "./ObservationSpec";
import { ObservationWriter } from "./ObservationWriter";

import * as PNG from "@vivaxy/png";
/**
 * @en
 * Sensor that wraps around another Sensor to provide temporal stacking.
 * Conceptually, consecutive observations are stored left-to-right, which is how they're output
 * For example, 4 stacked sets of observations would be output like
 *   |  t = now - 3  |  t = now -3  |  t = now - 2  |  t = now  |
 * Internally, a circular buffer of arrays is used. The m_CurrentIndex represents the most recent observation.
 * Currently, observations are stacked on the last dimension.
 * 
 * @zh
 * 传感器，包装另一个传感器以提供时间堆叠。
 * 从概念上讲，连续的观察结果是从左到右存储的，这就是它们的输出方式。
 * 例如，4个堆叠的观察结果集将被输出为
 *   |  t = now - 3  |  t = now -3  |  t = now - 2  |  t = now  |
 * 在内部，使用数组的循环缓冲区。m_CurrentIndex表示最近的观察结果。
 * 目前，观察结果是在最后一个维度上堆叠的。
 * 
 */
export class StackingSensor implements ISensor, IBuiltInSensor {

    private _wrappedSensor: ISensor = null!;
    private _numStackedObservations: number = 0;
    private _unstackedObservationSize: number = 0;
    private _name = "";
    private _observationSpec: ObservationSpec = null!;
    private _wrappedSpec: ObservationSpec = null!;

    private _stackedObservations: number[][];
    private _stackedCompressedObservations: Uint8Array[] = null!;

    private _currentIdx: number = 0;

    private _localWriter: ObservationWriter = null!;

    private _emptyCompressedObservation: Uint8Array = null!;
    private _compressionMapping: number[] = [];
    private _tensorShape: TensorShape = null!;

    constructor(wrappedSensor: ISensor, numStackedObservations: number) {

        this._wrappedSensor = wrappedSensor;
        this._numStackedObservations = numStackedObservations;
        this._name = `StackingSensor_size${numStackedObservations}_${wrappedSensor.getName()}`;

        this._wrappedSpec = wrappedSensor.getObservationSpec();

        this._unstackedObservationSize = SensorExtensions.observationSize(wrappedSensor);

        const newShape = this._wrappedSpec.shape;
        newShape[newShape.length - 1] *= numStackedObservations;

        this._observationSpec = new ObservationSpec(
            newShape, this._wrappedSpec.dimensionProperties, this._wrappedSpec.observationType);

        this._stackedObservations = new Array<number[]>(this._numStackedObservations);
        for (let i = 0; i < this._numStackedObservations; i++) {
            this._stackedObservations[i] = new Array<number>(this._unstackedObservationSize).fill(0);
        }

        if (this._wrappedSensor.getCompressionSpec().sensorCompressionType !== SensorCompressionType.None) {
            this._stackedCompressedObservations = new Array<Uint8Array>(this._numStackedObservations);
            this._emptyCompressedObservation = this.createEmptyPNG();
            for (let i = 0; i < this._numStackedObservations; i++) {
                this._stackedCompressedObservations[i] = this._emptyCompressedObservation;
            }
            this._compressionMapping = this.constructStackedCompressedChannelMapping(wrappedSensor);
        }

        if (this._wrappedSpec.rank !== 1) {
            const wrappedShape = this._wrappedSpec.shape;
            this._tensorShape = TensorShape.createWithArray([0, wrappedShape[0], wrappedShape[1], wrappedShape[2]]);
        }
    }

    public write(writer: ObservationWriter): number {
        const localWriter = new ObservationWriter();
        localWriter.setTarget(this._stackedObservations[this._currentIdx], this._wrappedSpec, 0);
        this._wrappedSensor.write(localWriter);

        let numWritten = 0;
        if (this._wrappedSpec.rank === 1) {
            for (let i = 0; i < this._numStackedObservations; i++) {
                const obsIndex = (this._currentIdx + 1 + i) % this._numStackedObservations;
                writer.addList(this._stackedObservations[obsIndex], numWritten);
                numWritten += this._unstackedObservationSize;
            }
        } else {
            for (let i = 0; i < this._numStackedObservations; i++) {
                const obsIndex = (this._currentIdx + 1 + i) % this._numStackedObservations;
                for (let h = 0; h < this._wrappedSpec.shape[0]; h++) {
                    for (let w = 0; w < this._wrappedSpec.shape[1]; w++) {
                        for (let c = 0; c < this._wrappedSpec.shape[2]; c++) {
                            writer.setItem3D(h, w, i * this._wrappedSpec.shape[2] + c, this._stackedObservations[obsIndex][this._tensorShape.index4D(0, h, w, c)]);
                        }
                    }
                }
            }
            numWritten = this._wrappedSpec.shape[0] * this._wrappedSpec.shape[1] * this._wrappedSpec.shape[2] * this._numStackedObservations;
        }

        return numWritten;
    }

    public update(): void {
        this._wrappedSensor.update();
        this._currentIdx = (this._currentIdx + 1) % this._numStackedObservations;
    }

    public reset(): void {
        this._wrappedSensor.reset();
        for (let i = 0; i < this._numStackedObservations; i++) {
            this._stackedObservations[i].fill(0);
        }
        if (this._wrappedSensor.getCompressionSpec().sensorCompressionType !== SensorCompressionType.None) {
            for (let i = 0; i < this._numStackedObservations; i++) {
                this._stackedCompressedObservations[i] = this._emptyCompressedObservation;
            }
        }
    }

    public getObservationSpec(): ObservationSpec {
        return this._observationSpec;
    }

    public getName(): string {
        return this._name;
    }

    public getCompressedObservation(): Uint8Array {
        const compressed = this._wrappedSensor.getCompressedObservation();
        this._stackedCompressedObservations[this._currentIdx] = compressed;
        let bytesLength = 0;
        for (const compressedObs of this._stackedCompressedObservations) {
            bytesLength += compressedObs.length;
        }
        const outputBytes = new Uint8Array(bytesLength);
        let offset = 0;
        for (let i = 0; i < this._numStackedObservations; i++) {
            const obsIndex = (this._currentIdx + 1 + i) % this._numStackedObservations;
            outputBytes.set(this._stackedCompressedObservations[obsIndex], offset);
            offset += this._stackedCompressedObservations[obsIndex].length;
        }
        return outputBytes;
    }

    public getCompressionSpec(): CompressionSpec {
        const wrappedSpec = this._wrappedSensor.getCompressionSpec();
        return new CompressionSpec(wrappedSpec.sensorCompressionType, this._compressionMapping);
    }

    private createEmptyPNG(): Uint8Array {
        const shape = this._wrappedSpec.shape;
        const height = shape[0];
        const width = shape[1];
        const numChannels = 4;

        const blackValueView = new Uint8Array(width * height * numChannels);
        for (let i = 0; i < width * height; i++) {
            const offset = i * numChannels;
            blackValueView[offset] = 0;
            blackValueView[offset + 1] = 0;
            blackValueView[offset + 2] = 0;
            blackValueView[offset + 3] = 0;

        }
        // const blackMemImageSource = {
        //     width: width,
        //     height: height,
        //     _data: blackValueView,
        //     _compressed: false,
        //     format: Texture2D.PixelFormat.RGBA8888,
        // };

        // const imgAsset = new ImageAsset(blackMemImageSource);
        // const blackTexture = new Texture2D();
        // blackTexture._uuid = 'black-texture';
        // blackTexture.image = imgAsset;


        const png = PNG.encode({
            width: width,
            height: height,
            data: Array.from(blackValueView),
            depth: 16,
            compression: 0,
            interlace: 0,
            colorType: 2,
            filter: 0,
            gamma: 1,
        });

        return png;
    }

    private constructStackedCompressedChannelMapping(wrappedSensor: ISensor): number[] {
        let wrappedMapping: number[] = [];
        const wrappedNumChannel = this._wrappedSensor.getObservationSpec().shape[2];
        wrappedMapping = wrappedSensor.getCompressionSpec().compressedChannelMapping ?? [];
        if (wrappedMapping.length === 0) {
            if (wrappedNumChannel === 1) {
                wrappedMapping = [0, 0, 0];
            } else {
                wrappedMapping = Array.from({ length: wrappedNumChannel }, (_, i) => i);
            }
        }

        const paddedMapLength = Math.ceil(wrappedMapping.length / 3) * 3;
        const compressionMapping = new Array<number>(paddedMapLength * this._numStackedObservations);
        for (let i = 0; i < this._numStackedObservations; i++) {
            const offset = wrappedNumChannel * i;
            for (let j = 0; j < paddedMapLength; j++) {
                if (j < wrappedMapping.length) {
                    compressionMapping[j + paddedMapLength * i] = wrappedMapping[j] >= 0 ? wrappedMapping[j] + offset : -1;
                } else {
                    compressionMapping[j + paddedMapLength * i] = -1;
                }
            }
        }
        return compressionMapping;
    }

    public getBuiltInSensorType(): BuiltInSensorType {
        const wrappedBuiltInSensor = this._wrappedSensor as unknown as IBuiltInSensor;
        return wrappedBuiltInSensor?.getBuiltInSensorType() ?? BuiltInSensorType.Unknown;
    }

    public getStackedObservations(): Array<number> {
        const observations: number[] = [];
        for (let i = 0; i < this._numStackedObservations; i++) {
            const obsIndex = (this._currentIdx + 1 + i) % this._numStackedObservations;
            observations.push(...this._stackedObservations[obsIndex]);
        }
        return observations;
    }
}
